{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "technological-removal",
   "metadata": {},
   "outputs": [],
   "source": [
    "{-# LANGUAGE DeriveAnyClass #-}\n",
    "{-# LANGUAGE DeriveGeneric #-}\n",
    "{-# LANGUAGE MultiParamTypeClasses #-}\n",
    "{-# LANGUAGE RecordWildCards #-}\n",
    "{-# LANGUAGE ScopedTypeVariables #-}\n",
    "\n",
    "import Control.Exception.Safe\n",
    "  ( SomeException (..),\n",
    "    try,\n",
    "  )\n",
    "import Control.Monad ( forM_, when, (<=<) )\n",
    "import Control.Monad.Cont ( ContT (..) )\n",
    "import GHC.Generics\n",
    "import Pipes hiding ( (~>) )\n",
    "import qualified Pipes.Prelude as P\n",
    "import Torch\n",
    "import Torch.Serialize\n",
    "import Torch.Typed.Vision ( initMnist )\n",
    "import qualified Torch.Vision as V\n",
    "import Prelude hiding ( exp )\n",
    "\n",
    "data MLP = MLP\n",
    "  { fc1 :: Linear,\n",
    "    fc2 :: Linear,\n",
    "    fc3 :: Linear\n",
    "  }\n",
    "  deriving (Generic, Show, Parameterized)\n",
    "\n",
    "data MLPSpec = MLPSpec\n",
    "  { i :: Int,\n",
    "    h1 :: Int,\n",
    "    h2 :: Int,\n",
    "    o :: Int\n",
    "  }\n",
    "  deriving (Show, Eq)\n",
    "\n",
    "\n",
    "(~>) :: (a -> b) -> (b -> c) -> a -> c\n",
    "f ~> g = g. f\n",
    "\n",
    "mlp :: MLP -> Tensor -> Tensor\n",
    "mlp MLP {..} = \n",
    "  -- Layer 1\n",
    "  linear fc1\n",
    "  ~> relu\n",
    "  \n",
    "  -- Layer 2\n",
    "  ~> linear fc2\n",
    "  ~> relu\n",
    "  \n",
    "  -- Layer 3\n",
    "  ~> linear fc3\n",
    "  ~> logSoftmax (Dim 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "robust-medium",
   "metadata": {},
   "outputs": [],
   "source": [
    "instance Randomizable MLPSpec MLP where\n",
    "  sample MLPSpec {..} =\n",
    "    MLP\n",
    "      <$> sample (LinearSpec i h1)\n",
    "      <*> sample (LinearSpec h1 h2)\n",
    "      <*> sample (LinearSpec h2 o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "single-representative",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainLoop :: Optimizer o => MLP -> o -> ListT IO (Tensor, Tensor) -> IO MLP\n",
    "trainLoop model optimizer = P.foldM step begin done. enumerateData\n",
    "  where\n",
    "    step :: MLP -> ((Tensor, Tensor), Int) -> IO MLP\n",
    "    step model ((input, label), iter) = do\n",
    "      let loss = nllLoss' label $ mlp model input\n",
    "      -- Print loss every 50 batches\n",
    "      when (iter `mod` 50 == 0) $ do\n",
    "        putStrLn $ \"Iteration: \" ++ show iter ++ \" | Loss: \" ++ show loss\n",
    "      (newParam, _) <- runStep model optimizer loss 1e-3\n",
    "      return newParam\n",
    "    done = pure\n",
    "    begin = pure model\n",
    "\n",
    "displayImages :: MLP -> (Tensor, Tensor) -> IO ()\n",
    "displayImages model (testImg, testLabel) = do\n",
    "  V.dispImage testImg\n",
    "  putStrLn $ \"Model        : \" ++ (show. argmax (Dim 1) RemoveDim. exp $ mlp model testImg)\n",
    "  putStrLn $ \"Ground Truth : \" ++ show testLabel\n",
    "\n",
    "main :: IO ()\n",
    "main = do\n",
    "  (trainData, testData) <- initMnist \"data\"\n",
    "  let trainMnist = V.MNIST {batchSize = 256, mnistData = trainData}\n",
    "      testMnist = V.MNIST {batchSize = 1, mnistData = testData}\n",
    "      spec = MLPSpec 784 64 32 10\n",
    "      optimizer = GD\n",
    "  net <- sample spec\n",
    "  \n",
    "  -- Train for 5 epochs\n",
    "  net' <- foldLoop net 5 $ \\model _ ->\n",
    "    runContT (streamFromMap (datasetOpts 2) trainMnist) $ trainLoop model optimizer. fst\n",
    "\n",
    "  -- Show test images + labels\n",
    "  forM_ [0 .. 10] $ displayImages net' <=< getItem testMnist\n",
    "\n",
    "  putStrLn \"Done\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "alien-mission",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Iteration: 0 | Loss: Tensor Float []  12.3775   \n",
       "Iteration: 50 | Loss: Tensor Float []  1.0952   \n",
       "Iteration: 100 | Loss: Tensor Float []  0.5626   \n",
       "Iteration: 150 | Loss: Tensor Float []  0.6660   \n",
       "Iteration: 200 | Loss: Tensor Float []  0.4771   \n",
       "Iteration: 0 | Loss: Tensor Float []  0.5012   \n",
       "Iteration: 50 | Loss: Tensor Float []  0.4058   \n",
       "Iteration: 100 | Loss: Tensor Float []  0.3095   \n",
       "Iteration: 150 | Loss: Tensor Float []  0.4237   \n",
       "Iteration: 200 | Loss: Tensor Float []  0.3433   \n",
       "Iteration: 0 | Loss: Tensor Float []  0.3671   \n",
       "Iteration: 50 | Loss: Tensor Float []  0.3206   \n",
       "Iteration: 100 | Loss: Tensor Float []  0.2467   \n",
       "Iteration: 150 | Loss: Tensor Float []  0.3420   \n",
       "Iteration: 200 | Loss: Tensor Float []  0.2737   \n",
       "Iteration: 0 | Loss: Tensor Float []  0.3054   \n",
       "Iteration: 50 | Loss: Tensor Float []  0.2779   \n",
       "Iteration: 100 | Loss: Tensor Float []  0.2161   \n",
       "Iteration: 150 | Loss: Tensor Float []  0.2933   \n",
       "Iteration: 200 | Loss: Tensor Float []  0.2289   \n",
       "Iteration: 0 | Loss: Tensor Float []  0.2693   \n",
       "Iteration: 50 | Loss: Tensor Float []  0.2530   \n",
       "Iteration: 100 | Loss: Tensor Float []  0.1979   \n",
       "Iteration: 150 | Loss: Tensor Float []  0.2616   \n",
       "Iteration: 200 | Loss: Tensor Float []  0.1986   \n",
       "              \n",
       "              \n",
       "              \n",
       "              \n",
       "   #%%*****   \n",
       "      ::: %   \n",
       "         %:   \n",
       "        :%    \n",
       "        #:    \n",
       "       :%     \n",
       "       %.     \n",
       "      #=      \n",
       "     :%.      \n",
       "     =#       \n",
       "Model        : Tensor Int64 [1] [ 7]\n",
       "Ground Truth : Tensor Int64 [1] [ 7]\n",
       "              \n",
       "              \n",
       "     %%%#     \n",
       "    %#  %     \n",
       "    .  #%     \n",
       "      :%:     \n",
       "      %+      \n",
       "     *%       \n",
       "     %=       \n",
       "    %%        \n",
       "    %%%%++%%%=\n",
       "     ==%%=.   \n",
       "              \n",
       "              \n",
       "Model        : Tensor Int64 [1] [ 2]\n",
       "Ground Truth : Tensor Int64 [1] [ 2]\n",
       "              \n",
       "              \n",
       "        .-    \n",
       "        =     \n",
       "        %     \n",
       "       .#     \n",
       "       =:     \n",
       "       @      \n",
       "       #      \n",
       "      ++      \n",
       "      %:      \n",
       "      %       \n",
       "              \n",
       "              \n",
       "Model        : Tensor Int64 [1] [ 1]\n",
       "Ground Truth : Tensor Int64 [1] [ 1]\n",
       "              \n",
       "              \n",
       "       %.     \n",
       "      *%-     \n",
       "     %%%%#    \n",
       "    :%%+:%-   \n",
       "    %%   -%.  \n",
       "    %    .@+  \n",
       "    %    %%.  \n",
       "    %   #%*   \n",
       "    %%%%%%    \n",
       "    :%%%-     \n",
       "              \n",
       "              \n",
       "Model        : Tensor Int64 [1] [ 0]\n",
       "Ground Truth : Tensor Int64 [1] [ 0]\n",
       "              \n",
       "              \n",
       "              \n",
       "     =    +   \n",
       "     %    %   \n",
       "    +.    %   \n",
       "    %    %:   \n",
       "    +    %    \n",
       "    %--=*%    \n",
       "     :: +%    \n",
       "        =%    \n",
       "        =%    \n",
       "        *     \n",
       "              \n",
       "Model        : Tensor Int64 [1] [ 4]\n",
       "Ground Truth : Tensor Int64 [1] [ 4]\n",
       "              \n",
       "              \n",
       "              \n",
       "        %@    \n",
       "        @:    \n",
       "       =@     \n",
       "       @%     \n",
       "       @      \n",
       "      :@      \n",
       "      %#      \n",
       "      @       \n",
       "      @       \n",
       "      +       \n",
       "              \n",
       "Model        : Tensor Int64 [1] [ 1]\n",
       "Ground Truth : Tensor Int64 [1] [ 1]\n",
       "              \n",
       "              \n",
       "              \n",
       "     %     %  \n",
       "    %     %   \n",
       "   +#    -+   \n",
       "   +%*::*%    \n",
       "    :%==%+    \n",
       "        %     \n",
       "       ++     \n",
       "       %      \n",
       "       %-+    \n",
       "       *      \n",
       "              \n",
       "Model        : Tensor Int64 [1] [ 4]\n",
       "Ground Truth : Tensor Int64 [1] [ 4]\n",
       "              \n",
       "              \n",
       "              \n",
       "      +       \n",
       "     %%+      \n",
       "    .%*%%     \n",
       "    -: *%     \n",
       "    -#-%%.    \n",
       "     %% =#    \n",
       "         %    \n",
       "         .%   \n",
       "          #.  \n",
       "           %  \n",
       "              \n",
       "Model        : Tensor Int64 [1] [ 9]\n",
       "Ground Truth : Tensor Int64 [1] [ 9]\n",
       "              \n",
       "              \n",
       "         ..=. \n",
       "      .%%%%%% \n",
       "     ::%+:    \n",
       "    %         \n",
       "   %          \n",
       "   %=         \n",
       "   %%%%%%+    \n",
       "     :%%%%    \n",
       "      %%%%    \n",
       "       %#     \n",
       "              \n",
       "              \n",
       "Model        : Tensor Int64 [1] [ 6]\n",
       "Ground Truth : Tensor Int64 [1] [ 5]\n",
       "              \n",
       "              \n",
       "              \n",
       "              \n",
       "      +%%%#   \n",
       "    +%*  .%%  \n",
       "   :%.  .#%+  \n",
       "    %@%%%%*   \n",
       "       +%-    \n",
       "      -%#     \n",
       "      %%      \n",
       "     %%       \n",
       "     %=       \n",
       "     @        \n",
       "Model        : Tensor Int64 [1] [ 9]\n",
       "Ground Truth : Tensor Int64 [1] [ 9]\n",
       "              \n",
       "              \n",
       "       ==:    \n",
       "     %%**%%   \n",
       "    .%    %:  \n",
       "    *-    +#  \n",
       "    %     :#  \n",
       "    #     :#  \n",
       "   -#     +#  \n",
       "   -#    .%   \n",
       "    #   +%:   \n",
       "    #%%%%=    \n",
       "              \n",
       "              \n",
       "Model        : Tensor Int64 [1] [ 0]\n",
       "Ground Truth : Tensor Int64 [1] [ 0]\n",
       "Done"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "main"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Haskell",
   "language": "haskell",
   "name": "haskell"
  },
  "language_info": {
   "codemirror_mode": "ihaskell",
   "file_extension": ".hs",
   "mimetype": "text/x-haskell",
   "name": "haskell",
   "pygments_lexer": "Haskell",
   "version": "8.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
